import requests
import numpy as np
from typing import ClassVar
from config_load import app_config
from langchain_openai import ChatOpenAI,OpenAIEmbeddings
from langchain_ollama import OllamaEmbeddings


class SILICONFLOWRerank():
    """
        初始化rerank模型
    """
    RANK_MODEL: ClassVar[str]       = app_config['siliconflow']['rank_model']
    Rank_URL: ClassVar[str]         = app_config['siliconflow']['rank_url']
    API_KEY: ClassVar[str]          = app_config['siliconflow']['api_key']
    
    
    def __init__(self) -> None:

        base_url = self.Rank_URL
        self.model_name = self.RANK_MODEL
        self.base_url = base_url
        
        self.headers = {
            "accept": "application/json",
            "content-type": "application/json",
            "authorization": f"Bearer {self.API_KEY}",
        }

    def similarity(self, query: str, texts: list):
        
        payload = {
            "model": self.model_name,
            "query": query,
            "documents": texts,
            "top_n": len(texts),
            "return_documents": False,
            "max_chunks_per_doc": 1024,
            "overlap_tokens": 80,
        }
        
        response = requests.post(
            self.base_url, json=payload, headers=self.headers
        ).json()
        
        rank = np.zeros(len(texts), dtype=float)
        
        if "results" not in response:
            return rank, 0

        for d in response["results"]:
            rank[d["index"]] = d["relevance_score"]
        return (
            rank,
            response["meta"]["tokens"]["input_tokens"] + response["meta"]["tokens"]["output_tokens"],
        )





class OllamaRerank:
    """基于Ollama的嵌入式重排序模型"""
    def __init__(self, model_name: str = "llama3", temperature: float = 0.2):
        self.model_name = model_name
        self.temperature = temperature
        self.embedder = OllamaEmbeddings(model=model_name)
        
    def _generate_embeddings(self, texts: list) -> list:
        """批量生成Ollama嵌入向量"""
        return [self.embedder.embed(text) for text in texts]

    def rerank(self, query: str, candidates: list, top_k: int = 3) -> list:
        """
        Ollama驱动的重排序逻辑
        Args:
            query: 原始查询
            candidates: 候选列表（含原始文本和得分）
            top_k: 返回结果数量
        Returns:
            重排序后的候选列表
        """
        # 生成查询嵌入
        query_embed = self.embedder.embed([query])[0]
        
        # 批量生成候选嵌入
        candidate_texts = [c['text'] for c in candidates]
        candidate_embeds = self._generate_embeddings(candidate_texts)
        
        # 计算相似度得分
        scores = []
        for ce, cc in zip(candidate_embeds, candidates):
            # 使用Ollama的语义相似度计算
            sim = self._cosine_similarity(query_embed, ce)
            # 混合原始得分和语义得分
            hybrid_score = 0.6 * cc['score'] + 0.4 * sim
            scores.append(hybrid_score)
        
        # 排序并返回top_k
        ranked = sorted(zip(candidate_texts, scores), 
                       key=lambda x: -x[1])[:top_k]
        return [{
            'text': text,
            'score': score,
            'metadata': next((c['metadata'] for c in candidates 
                            if c['text'] == text), {})
        } for text, score in ranked]

    @staticmethod
    def _cosine_similarity(a: list, b: list) -> float:
        """余弦相似度计算"""
        import numpy as np
        a = np.array(a).reshape(1, -1)
        b = np.array(b).reshape(1, -1)
        return np.dot(a, b.T) / (np.linalg.norm(a) * np.linalg.norm(b))




    

            

    